// 使用mongodb的GridFS来存储上传的文件

package storage

import (
	"io/ioutil"

	"gopkg.in/mgo.v2"
)

type (
	// MongoDBStorage struct
	MongoDBStorage struct {
		MgoURL string // mongodb的连接url
		Prefix string // GridFS的
	}
)

var originalSession *mgo.Session

func init() {
	Register("mongodb", NewMongoDBStorage())
}

// NewMongoDBStorage function
func NewMongoDBStorage() *MongoDBStorage {
	return &MongoDBStorage{}
}

// Get function
// 获得图片
func (ms *MongoDBStorage) Get(src string) ([]byte, error) {
	session := getSession(ms.MgoURL)
	defer session.Close()
	gridfs := ms.getGridFS(session)

	fd, err := gridfs.Open(src)
	if err != nil {
		return nil, err
	}
	defer fd.Close()
	imageData, err := ioutil.ReadAll(fd)
	if err != nil {
		return nil, err
	}
	return imageData, nil
}

// Save function
// 存储图片
func (ms *MongoDBStorage) Save(src string, imageData []byte) error {
	session := getSession(ms.MgoURL)
	defer session.Close()
	gridfs := ms.getGridFS(session)
	fd, err := gridfs.Create(src)
	if err != nil {
		return err
	}
	defer fd.Close()

	_, err = fd.Write(imageData)
	return err
}

// Delete Funtion
// 删除图片
func (ms *MongoDBStorage) Delete(src string) error {
	session := getSession(ms.MgoURL)
	defer session.Close()

	return ms.getGridFS(session).Remove(src)
}

// IsExist Funtion
// 检查文件是否存在
func (ms *MongoDBStorage) IsExist(src string) bool {
	session := getSession(ms.MgoURL)
	defer session.Close()
	gridfs := ms.getGridFS(session)

	fd, err := gridfs.Open(src)
	if err != nil {
		// mgo.ErrNotFound
		return false
	}
	defer fd.Close()
	return true
}

func (ms *MongoDBStorage) getGridFS(session *mgo.Session) *mgo.GridFS {
	return session.DB("").GridFS(ms.Prefix)
}

func dial(url string) *mgo.Session {
	session, err := mgo.Dial(url)
	if err != nil {
		panic(err)
	}
	return session
}

func getSession(url string) *mgo.Session {
	if originalSession == nil {
		originalSession = dial(url)
	}
	return originalSession.New()
}
